import json
import os

# 路径根据实际情况修改
base_dir = "/home/ai999/project/A2XP-master/data/ThyroidPACS"
domain_json = os.path.join(base_dir, "domain_names.json")
train_txt = os.path.join(base_dir, "thyroid_train.txt")
test_txt  = os.path.join(base_dir, "thyroid_test.txt")

# 输出新文件
train_out = os.path.join(base_dir, "thyroid_train_with_domain.txt")
test_out  = os.path.join(base_dir, "thyroid_test_with_domain.txt")

# 加载域名映射
with open(domain_json, "r", encoding="utf-8") as f:
    domain_map = json.load(f)

# 我们只需要 value 集合（医院名列表）
valid_domains = set(domain_map.values())

def add_domain(infile, outfile):
    with open(infile, "r", encoding="utf-8") as fin, \
         open(outfile, "w", encoding="utf-8") as fout:
        for line in fin:
            line = line.strip()
            if not line:
                continue
            parts = line.split()
            if len(parts) != 2:
                raise ValueError(f"行格式不符合预期: {line}")
            rel_path, label = parts

            # 域名从路径的第一段拿，例如 "301桥本结节/xxx/yyy.jpg"
            domain_name = rel_path.split("/")[0]
            if domain_name not in valid_domains:
                raise ValueError(f"域名 {domain_name} 不在 domain_names.json 里")

            fout.write(f"{rel_path} {label} {domain_name}\n")

# 处理 train/test
add_domain(train_txt, train_out)
add_domain(test_txt, test_out)

print("转换完成！生成文件：")
print("  ", train_out)
print("  ", test_out)
